#include "clipp.h"
#include "dumpParser.h"
#include <cstdlib>
#include <fstream>
#include <graphviz/gvc.h>
#include <iostream>
#include <string>
#include <sys/stat.h>
#include <unordered_set>
#include <regex>

using namespace clipp;
using namespace std;

bool noplt = false;
bool nostd = false;
void printHelpMessage(string programName) {
  cout << "Usage: " << programName << " <option(s)> <file>" << endl;
  cout << " Extracting Function Call Graph from [Binary/asm] <Files>" << endl;
  cout << " At least one of the following switches must be given:" << endl;
  cout << "  -h, --help                  Print this help message" << endl;
  cout << "  -o, --output <file>         Output the call graph to <file>"
       << endl;
  cout << " The following switches are optional:" << endl;
  // Add more commands here
}
bool isObjdumpAvailable() {
  // Check if objdump command is available by trying to execute it with the
  // --version option
  int result = system("objdump --version > /dev/null 2>&1");
  return result == 0;
}
bool dumpFile(const std::string &filename, const std::string &dumpFilename) {
  struct stat buffer;
  if (stat(filename.c_str(), &buffer)) {
    std::cout << "[error] file not found: " << filename << std::endl;
    return false;
  }
  if (buffer.st_size == 0) {
    std::cout << "[error] file is empty: " << filename << std::endl;
    return false;
  }
  if (!isObjdumpAvailable()) {
    cout << "[error] objdump command not found" << std::endl;
    return false;
  }

  std::string command = "objdump -dS " + filename + " > " + dumpFilename;
  int result = system(command.c_str());
  if (result != 0) {
    std::cout << "[error] objdump failed" << std::endl;
    return false;
  }
  return true;
}
std::string demangle(const std::string mangled_name) {
  std::string command = "c++filt \"" + mangled_name + "\"";
  std::string result;
  FILE *pipe = popen(command.c_str(), "r");
  if (!pipe) {
    std::cerr << "popen failed!" << std::endl;
    return "";
  }
  char buffer[128];
  while (!feof(pipe)) {
    if (fgets(buffer, 128, pipe) != nullptr) {
      result += buffer;
    }
  }
  pclose(pipe);
  if (result.back() == '\n')
    result.pop_back();
  return result;
}
string getTrimmedFunctionName(string result, long unsigned int maxLength) {
  if (result.size() < maxLength)
    return result;

  size_t pos = result.find("(");
  if (pos != std::string::npos) {
    result = result.substr(0, pos);
    result += "(..)";
  }
  
  pos = result.rfind("::");
  if (pos != std::string::npos && result.size() > maxLength) {
    string endStr = result.substr(pos, result.size());
    size_t len = maxLength - (result.size() - pos);
    string startStr = result.substr(0, len > maxLength ? 0:len );
    result = startStr + (len > maxLength ? "":"..") + endStr;
  }
  if(result.size() > maxLength + 4)
    result = result.substr(0, maxLength - 4) + "(..)";

  // Replace all '&' with '&amp;'
  pos = result.find("&");
  while (pos != std::string::npos) {
    result.replace(pos, 1, "&amp;");
    pos = result.find("&", pos + 4);
  }
  // Replace all '<' with '<'
  pos = result.find("<");
  while (pos != std::string::npos) {
    result.replace(pos, 1, "&lt;");
    pos = result.find("<", pos + 4);
  }

  // Replace all ">" with "&gt;"
  pos = result.find(">");
  while (pos != std::string::npos) {
    result.replace(pos, 1, "&gt;");
    pos = result.find(">", pos + 4);
  }
  pos = result.find("\"");
  while (pos != std::string::npos) {
    result.replace(pos, 1, "&quot;");
    pos = result.find("\"", pos + 6);
  }
  pos = result.find("'");
  while (pos != std::string::npos) {
    result.replace(pos, 1, "&apos;");
    pos = result.find("'", pos + 6);
  }
  // pos = result.find(" ");
  // while (pos != std::string::npos) {
  //   result.replace(pos, 1, "&nbsp;");
  //   pos = result.find(" ", pos + 5);
  // }
  return result;
}
bool matchSkipPatten(string funcName, string skipFunc) {
  if(skipFunc.empty())
    return false;
  std::regex reg(skipFunc);
  return std::regex_search(funcName, reg);
}
bool getDotFile(unordered_map<string, unordered_set<string>> &funcMap,
                const std::string &startFunc_, const std::string &dotFilename,
                long unsigned int maxLength,string skipFunc) {
  std::ofstream outfs(dotFilename.c_str());
  if (!outfs.is_open()) {
    std::cout << "[error] can not open file:" << dotFilename << std::endl;
    return false;
  }
  outfs << "digraph \"" << dotFilename << "\" {\n"
        << "node [fontname=courier, shape=box, style=filled, "
           "colorscheme=brbg9]\n"
           "rankdir=LR\n";
  vector<string> vStartFuncSet{startFunc_};
  unordered_set<string> sVisitedSet{startFunc_};
  vector<string> vplt;
  vector<string> vstd;
  while (!vStartFuncSet.empty()) {
    string startFunc = vStartFuncSet.back();
    vStartFuncSet.pop_back();
    string startFuncA = getTrimmedFunctionName(demangle(startFunc), maxLength);
    for (const auto &funcTo : funcMap[startFunc]) {
      if (funcTo.find("@plt") != std::string::npos) {
        vplt.push_back(getTrimmedFunctionName(
            demangle(funcTo.substr(0, funcTo.size() - 4)), maxLength));
        continue;
      }
      string funcName = demangle(funcTo);
      if(matchSkipPatten(funcName,skipFunc)) {
        sVisitedSet.insert(funcTo);
        continue;
      }
        
      if (funcName.find("std::") == 0 || funcTo.find("_ZNSt6") == 0) {
        string result = getTrimmedFunctionName(funcName, maxLength);
        vstd.push_back(result);
        continue;
      }
      if (!sVisitedSet.count(funcTo)) {
        vStartFuncSet.push_back(funcTo);
        sVisitedSet.insert(funcTo);
      }
      outfs << "\"" << startFuncA << "\" -> \""
            << getTrimmedFunctionName(funcName, maxLength) << "\""
            << ";" << endl;
    }
  }
  if (!noplt && !vplt.empty()) {
    outfs << "\"@plt\" [label=<\"" << vplt[0] << "\"";
    for (size_t i = 1; i < vplt.size(); i++) {
      outfs << "<br/>\"" << vplt[i] << "\"";
    }
    outfs << ">; color=orange];\n\"" << startFunc_
          << "\" -> \"@plt\" [label=<@plt>];\n";
  }
  if (!nostd && !vstd.empty()) {
    outfs << "\"std::\" [label=<\"" << vstd[0] << "\"";
    for (size_t i = 1; i < vstd.size(); i++) {
      outfs << "<br/>\"" << vstd[i] << "\"";
    }
    outfs << ">; color=orange];\n\"" << startFunc_
          << "\" -> \"std::\" [label=<std::>];\n";
  }

  outfs << "}\n";
  outfs.close();
  return true;
}
int main(int argc, char *argv[]) {
  string filename, startFunc;
  string dotFilename;
  bool helpFlag = false;
  long unsigned int maxFuncLen = 36;
  string skipFunc;
  auto cli =
      (value("input file", filename),
       option("-h", "-H", "--help").set(helpFlag),
       option("--noplt").set(noplt),
       option("--nostd").set(nostd),
       option("-s", "--startFunc") & value("startFunc", startFunc),
       option("-S", "--skip") & value("skipFunc", skipFunc),
       option("-m", "--maxFuncLen") & value("maxFuncLen", maxFuncLen),
       option("-o", "-O", "--dotFilename") & value("dotFilename", dotFilename));
  if (!parse(argc, const_cast<char **>(argv), cli) || helpFlag) {
    printHelpMessage(argv[0]);
    exit(1);
  }

  cout << "[message] filename: " << filename << "\n"
       << "[message] skipFunc: " << skipFunc << "\n"
       << "[message] startFunc: " << startFunc << endl;
  // if (!matchSkipPatten("clipp::detail::match_t::~match_t()", skipFunc)) {
  //   cout << "[error] can not match function" << endl;
  // }else {
  //   cout << "[message] okk" << endl;
  // }
  // exit(1);
  string dumpFilename = filename + ".dump";
  if (!dumpFile(filename, dumpFilename)) {
    cout << "[error] Please provide the binary file or the file after objdump."
         << endl;
    exit(1);
  }

  DumpParser dumpParser(dumpFilename);
  unordered_map<string, unordered_set<string>> funcMap =
      std::move(dumpParser.parse());
  // dumpParser.print(funcMap, startFunc);

  if (dotFilename.empty())
    dotFilename = startFunc + ".dot";
  if (!getDotFile(funcMap, startFunc, dotFilename, maxFuncLen,skipFunc)) {
    cout << "[error] dot file generation failed." << endl;
    return -1;
  }

  // dot command transfor
  std::string command =
      "dot -Tsvg " + dotFilename + " -o " + dotFilename + ".svg";
  int result = system(command.c_str());
  if (result != 0) {
    cout << "[error] dot command failed." << endl;
    return -1;
  }
  return 0;
}
